"""
将 icpr2012 数据集转换成 yolo 训练集
http://ludo17.free.fr/mitos_2012/index.html

"""
import csv
import glob
import os
import numpy as np
from PIL import Image
import cv2
from tqdm import tqdm

if __name__ == '__main__':
    wsi_path = '/media/hsmy/wanghao_18T/dataset/ICPR2012/A00_v2/'
    png_path = '/media/hsmy/wanghao_18T/dataset/ICPR2012/yolo/images/'
    label_path = '/media/hsmy/wanghao_18T/dataset/ICPR2012/yolo/labels/'
    mask_path = '/media/hsmy/wanghao_18T/dataset/ICPR2012/yolo/masks/'
    os.makedirs(mask_path, exist_ok=True)
    os.makedirs(png_path, exist_ok=True)
    os.makedirs(label_path, exist_ok=True)

    patch_size = 1024  # 训练 image 大小

    for file_path in tqdm(glob.glob(os.path.join(wsi_path, '*.bmp'))):
        file_name = file_path.split('.')[0]
        file_name = file_name.split('/')[-1]

        img = cv2.imread(file_path)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        width, height, _ = img.shape
        mask = np.zeros((height, width)).astype(np.uint8)

        csv_path = os.path.join(wsi_path, file_name + '.csv')
        csv_reader = csv.reader(open(csv_path, 'r'))
        coordinate_arr = []
        for row in csv_reader:
            x = np.array(list(map(int, row[::2]))).mean()
            y = np.array(list(map(int, row[1::2]))).mean()
            point = [int(x), int(y)]
            coordinate_arr.append(point)
            mask[int(y), int(x)] = 1

        index = 0  # 平移切割patch形成随机点位
        react_size_percent = 64 / patch_size
        for y in range(0, height, patch_size):
            if y + patch_size > height:
                y = height - patch_size  # 保证patch规格，超出边界往前推
            for x in range(0, width, patch_size):
                if x + patch_size > width:
                    x = width - patch_size
                # 定义块的范围
                y_end = min(y + patch_size, height)
                x_end = min(x + patch_size, width)

                # 切割块
                png_patch = img[y:y_end, x:x_end]
                mask_patch = mask[y:y_end, x:x_end]

                if mask_patch.max() == 0:
                    continue

                png_name = f"{file_name}_{index}.png"
                txt_name = f"{file_name}_{index}.txt"

                points = np.argwhere(mask_patch == 1)
                for point in points:
                    label = (
                        0,
                        point[1] / patch_size,
                        point[0] / patch_size,
                        react_size_percent,
                        react_size_percent
                    )
                    with open(label_path + txt_name, 'a') as f:
                        f.write(('%g ' * len(label)).rstrip() % label + '\n')

                Image.fromarray(png_patch).save(png_path + png_name)

                kernel = np.ones((10, 10), np.uint8)  # 定义膨胀结构元素
                diffused_matrix = cv2.dilate(mask_patch, kernel, iterations=1)
                Image.fromarray((diffused_matrix * 255).astype(np.uint8)).save(mask_path + png_name)
                index += 1
